/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \file

    Contains the dl07_template_grain_model. */

// $Id$

#ifndef __dl07_template_grain_model__
#define __dl07_template_grain_model__

#include <vector>
#include "grain_collection_grain_model.h"
#include "grain_size.h"
#include "thermal_equilibrium_grain.h"
#include "misc.h"
#include "CCfits/CCfits"
#include "fits-utilities.h"
#include "blitz-fits.h"
#include "vecops.h"

namespace mcrx {
  template <template<class> class, typename> class dl07_template_grain_model;
  template <template<class> class chromatic_policy, typename T_rng_policy> 
  class dl07_template_grain_model_creator_functor;
};


/** This class implements the Draine & Li 07 grain model which is like
    the wd01_grain_model except that it uses the precomputed emission
    spectra from Drain & Li 07 as a function of local radiation
    intensity to calculate the dust emission.  */
template <template<class> class chromatic_policy, 
	  typename T_rng_policy = mcrx::global_random> 
class mcrx::dl07_template_grain_model : 
  public mcrx::grain_collection_grain_model<chromatic_policy, T_rng_policy> {

  /** The emission templates indexed by (intensity, wavelength). These
      are normalized to unit L_bol. */
  array_2 emission_sed_;
  /// The values of the intensity for the template entries.
  std::vector<T_float> U_;
  /// The midpoint values of the intensity, for searching.
  std::vector<T_float> Umid_;
  /// Wavelength vector of the emission template
  array_1 elambda_;
  /// Total absorption cross section for one mass unit of dust.
  array_1 sigma_abs_;
  /// The minimum value of U to use. (Dust with lower U uses Umin_,
  /// while enforcing energy conservation.)
  T_float Umin_;

public:
  dl07_template_grain_model (const Preferences& p, const T_unit_map&);
  dl07_template_grain_model (const dl07_template_grain_model& g) :
    // need to make sure array members are copied. We make the copies
    // thread-local, since the methods are not reentrant anyway.
    grain_collection_grain_model<chromatic_policy, T_rng_policy>(g),
    emission_sed_(g.emission_sed_.copy()), U_(g.U_), Umid_(g.Umid_),
    elambda_(g.elambda_.copy()), 
    sigma_abs_(g.sigma_abs_.copy()), 
    Umin_(g.Umin_)
  {};

  using mcrx::grain_collection_grain_model<chromatic_policy, T_rng_policy>::alambda;

  virtual boost::shared_ptr<grain_model<chromatic_policy, T_rng_policy> > 
  clone() const {
    return boost::shared_ptr<grain_model<chromatic_policy, T_rng_policy> >
      (new dl07_template_grain_model(*this));};

  /** Redefinition of the resample function because this class needs
      to update its own members when the wavelengths are resampled. */
  virtual void resample (const array_1& lambda, const array_1& elambda=array_1());

  template<typename T>
  void calculate_SED_virtual(const blitz::ETBase<T>& intensity, 
			     const array_1& m_dust,
			     array_2& sed, array_2* temp) const;
};



template <template<class> class chromatic_policy, 
	  typename T_rng_policy> 
mcrx::dl07_template_grain_model<chromatic_policy, T_rng_policy>::
dl07_template_grain_model (const Preferences& p,
		  const T_unit_map& units) :
  grain_collection_grain_model<chromatic_policy, T_rng_policy>()
{
  using namespace CCfits;
  using namespace blitz;

  const std::string data_directory =
      word_expand(p.getValue("grain_data_directory", std::string ()))[0];

  // shortcuts
  std::vector<boost::shared_ptr<emitting_grain> >& grains = 
    grain_collection_grain_model<chromatic_policy, T_rng_policy>::grains_;
  std::vector<array_1>& dnda =
    grain_collection_grain_model<chromatic_policy, T_rng_policy>::dnda_;

  // copy units from those supplied into *scatterer* unit map.
  T_unit_map& u = 
    mcrx::scatterer<chromatic_policy, T_rng_policy>::units_;
  u["length"] = units.get("length");
  u["wavelength"] = units.get("wavelength");
  u["mass"] = units.get("mass");

  // The grains use their preferred units, which are also copied into
  // the grain model units.

  // the grains are resampled onto a common grid, which is the first
  // grain's wavelengths. Since the PAHs have higher-resolution grids,
  // we put those first.
  grains.push_back
    (boost::shared_ptr<emitting_grain>
     (new thermal_equilibrium_grain
      (data_directory + "/PAH_neutral_DL07.fits", p)));
  grains.push_back
    (boost::shared_ptr<emitting_grain>
     (new thermal_equilibrium_grain
      (data_directory + "/PAH_ionized_DL07.fits", p)));
  // we must make sure we cut off the graphite grains at the largest
  // size of the pah's
  grains.push_back
    (boost::shared_ptr<emitting_grain>
     (new thermal_equilibrium_grain
      (data_directory + "/graphite.fits", p, 
       grains[0]->sizes().back()*1.00001)));
  grains.push_back
    (boost::shared_ptr<emitting_grain>
     (new thermal_equilibrium_grain
      (data_directory + "/suv_silicate.fits", p)));

  T_unit_map& u2 = 
    mcrx::grain_model<chromatic_policy, T_rng_policy>::units_;
  u2["length"]=grains.front()->units().get("length");
  u2["wavelength"]=grains.front()->units().get("wavelength");
  // grains have no mass unit. But the size distributions do, so for
  // simplicity we select the scatterer mass unit as our mass unit.
  // NOTE that this arrangement means the size distributions and the
  // grains in general use different length units.
  u2["mass"] = u.get("mass");

  //  size distributions are made with our set *scatterer* units
  const std::string wd01_set =
    p.getValue("wd01_parameter_set", std::string());
  wd01_graphite_distribution graphite_size_distribution(wd01_set, units);
  wd01_silicate_distribution silicate_size_distribution(wd01_set, units);

  dnda.resize(grains.size());
  dnda[0].resize(grains[0]->asizes().size());
  dnda[1].resize(grains[1]->asizes().size());
  dnda[2].resize(grains[2]->asizes().size());
  dnda[3].resize(grains[3]->asizes().size());

  // evaluate the size distributions. Note that we need to convert
  // from grain length unit here to the scatterer units
  const T_float gsize_conv = units::convert(u2.get("length"),
					    u.get("length"));
  // ionized fraction is a simple fit to Fig 8 in DL07.
  const array_1 x_ion(0.55*log10(units::convert(u2.get("length"),"angstrom")
				 *grains[0]->asizes())-0.1);
  ASSERT_ALL(x_ion>=0);
  ASSERT_ALL(x_ion<=1);
  dnda[0] = (1-x_ion)*
    graphite_size_distribution.dn_da(grains[0]->asizes()*gsize_conv);
  dnda[1] = x_ion*
    graphite_size_distribution.dn_da(grains[1]->asizes()*gsize_conv);
  dnda[2] = graphite_size_distribution.dn_da(grains[2]->asizes()*gsize_conv);
  dnda[3] = silicate_size_distribution.dn_da(grains[3]->asizes()*gsize_conv);
  ASSERT_ALL(dnda[0]==dnda[0]);
  ASSERT_ALL(dnda[1]==dnda[1]);
  ASSERT_ALL(dnda[2]==dnda[2]);
  ASSERT_ALL(dnda[3]==dnda[3]);


  // we also need to load the emission templates
  // templates are in subdir "DL07_emission" of data dir
  const string template_fn(data_directory + "/DL07_emission/" + 
			   wd01_set+".fits");
  std::cout << "Reading Draine & Li 07 emission templates from " 
	    << template_fn << std::endl;
  FITS template_file(template_fn);
  ExtHDU& sed_hdu=open_HDU(template_file, "SED");
  read(sed_hdu.column("L_lambda"), emission_sed_);
  sed_hdu.column("intensity").read(U_,1,emission_sed_.extent(firstDim));
  ExtHDU& lambda_hdu=open_HDU(template_file, "LAMBDA");
  read(lambda_hdu.column("lambda"),elambda_);
  // templates end at 1um. "extend" them in a very crude way
  elambda_(0)=alambda()(0);
  emission_sed_(Range::all(), Range(0,2)) = 0;

  // Calculate midpoints
  Umid_ =calculate_middle(U_);

  // Get Umin value, if defined
  Umin_ = p.defined("dl07_template_umin") ?
    p.getValue("dl07_template_umin", T_float()) : -blitz::huge(T_float());
  if(p.defined("dl07_template_umin")) {
    std::cout << "\tMinimum intensity template: " << Umin_ << std::endl;
  }
}


template <template<class> class chromatic_policy, 
	  typename T_rng_policy> 
void
mcrx::dl07_template_grain_model<chromatic_policy, T_rng_policy>::
resample (const array_1& lambda, const array_1& el)
{
  using namespace blitz;

  if((lambda.size()==alambda().size()) && (el.size()==elambda_.size()) && 
     all(lambda==alambda()) && all(el==elambda_))
    // nothing to do
    return;

  // first thing to do is to set up the base class
  grain_collection_grain_model<chromatic_policy, T_rng_policy>::
    resample(lambda, el);

  const array_1 elambda((el.size()==0)?lambda:el);

  // check if we are upsampling or downsampling. If upsampling, we use
  // mcrx::resample with a logarithmic resampling, because that works
  // better. For downsampling, we use mcrx::subsample, which preserves
  // the integral. Since the emission grid is larger than the
  // absorption, we use the emission grid to determine this
  const bool upsample = 
    (log10(alambda()(alambda().size()-1))-log10(alambda()(0)))/(alambda().size()-1) >
    ((log10(elambda(elambda.size()-1))-log10(elambda(0)))/
     (elambda.size()-1));
  if (upsample) 
    std::cout << "Upsampling DL07 templates\n";
  else
    std::cout << "Downsampling DL07 templates\n";

  // now that the base classes are updated for new wavelengths, we
  // need to do it for our members
  array_2 new_sed(shape(U_.size(), elambda.size()));
  for (int i=0; i< U_.size(); ++i) {
    if(upsample)
      new_sed(i, Range::all()) = 
	mcrx::resample(emission_sed_(i, Range::all()), elambda_, elambda,
		       false, true);
    else
      new_sed(i, Range::all()) = 
	mcrx::subsample(emission_sed_(i, Range::all()), elambda_, elambda);

    // renormalize to unit L_bol
    const double new_L_bol = integrate_quantity(new_sed(i, Range::all()), 
						elambda, false);
    new_sed(i, Range::all())/= new_L_bol;
  }
  emission_sed_.reference(new_sed);
  elambda_.reference(elambda.copy());

  // calculate total cross section for all grains
  sigma_abs_.resize(alambda().size());
  sigma_abs_ = this->sigma_tot(1.0);
}

template <template<class> class chromatic_policy, 
	  typename T_rng_policy> 
template<typename T>
void
mcrx::dl07_template_grain_model<chromatic_policy, T_rng_policy>::
calculate_SED_virtual(const blitz::ETBase<T>& intens, const array_1& m_dust,
		      array_2& sed, array_2* temp) const
{
  using namespace blitz;
  const T& intensity(intens.unwrap());

  if(temp)
    // temp does not make sense for this grain
    *temp = blitz::quiet_NaN(T_float());

  // simple loop over all cells
  const int cmax=m_dust.size();
  for (int c=0; c<cmax; ++c) {

    // Calculate radiative intensity. 8.65e-14 J/m3 is the definition of 1U
    const T_float U = 
      std::max(integrate_quantity(intensity(c,Range::all()),
				  alambda(), false)*
	       4*constants::pi/(constants::c0*8.65e-14),
	       Umin_);
    
    // find the sed index
    const int i = lower_bound (Umid_.begin(), Umid_.end(), U) -
      Umid_.begin() - 1;
    assert (i >= 0);
    assert (i < emission_sed_.extent(firstDim));

    // calculate the bolometric absorption luminosity so we can scale
    // the template

    const T_float L_bol = 
      integrate_quantity(sigma_abs_* intensity(c,Range::all()),
			 alambda(),
			 false) * 
      m_dust(c) * 4*constants::pi*this->sed_unit_conversion_;

    DEBUG(1,std::cout << "U= " << U << ", using template " << i << " for " << U_[i] << "\nL_bol = "<< L_bol << " m_dust = " << m_dust(c) << std::endl;);
    
    sed(c, Range::all()) += emission_sed_(i, Range::all())*L_bol;
  }
}



template <template<class> class chromatic_policy, typename T_rng_policy> 
class mcrx::dl07_template_grain_model_creator_functor
{
public:
  typedef boost::shared_ptr<mcrx::grain_model<chromatic_policy, T_rng_policy> > return_type;
  
  dl07_template_grain_model_creator_functor() {};
  return_type operator()(const Preferences& p) const {

    //THIS IS A HACK BECAUSE WE CAN ONLY HAVE ONE ARGUMENT
    T_unit_map units;
    units["mass"] = p.getValue("massunit", std::string());
    units["length"] = p.getValue("lengthunit", std::string());
    units["wavelength"] = p.getValue("wavelengthunit", std::string());

    return return_type
      ( new mcrx::dl07_template_grain_model<chromatic_policy, T_rng_policy> 
	(p, units));
  };
};


#endif
